from flask import Flask, render_template, request, redirect, url_for, jsonify, session
from flask_sqlalchemy import SQLAlchemy
from base64 import b64encode
import sqlalchemy as sa
import random
import traceback
import time
import json
import re
import os
from typing import Any

from .lang import random_gb2312_name
from .tts import play_audio_text
from .repl import pretty_print_message, pretty_print_embedding
from .mark import parse_md_page
from .core import ask_question
from .trunc import conversation_title, conversation_choices
from .safety import get_safety_timeout
from .embed import auto_embed_messages
from .config import get_option, get_options

app = Flask(__name__)
app.secret_key = os.urandom(24)
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///users.db'  # 使用 SQLite 数据库，保存在当前目录下的 users.db 文件
# app.config["SQLALCHEMY_ECHO"] = True
app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
db: Any = SQLAlchemy(app)

class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(32), unique=True, nullable=False)
    password = db.Column(db.String(32), nullable=False)
    email = db.Column(db.String(100), unique=True, nullable=False)
    admin = db.Column(db.Boolean, default=False)

    def __repr__(self):
        return f'User(username={self.username}, password={self.password}, admin={self.admin})'

class Robot(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    owner = db.Column(db.Integer, sa.ForeignKey('user.id'), nullable=False)
    botname = db.Column(db.String(32), unique=True, nullable=False)
    model = db.Column(db.Enum('gpt-4o', 'gpt-4o-mini', 'gpt-3.5-turbo', 'gpt-4-turbo', 'gpt-4-vision'), default='gpt-4o')
    temperature = db.Column(db.Float(), nullable=False, default=0.5)
    max_tokens = db.Column(db.Integer(), nullable=False, default=4096)
    introduction = db.Column(db.Text(), default=None)
    instruction = db.Column(db.Text(), default=None)
    tools = db.Column(db.Text(), default=' '.join(get_option('tool_categories', [])))

    def __repr__(self):
        return f'Robot(id={self.id}, owner={self.owner}, botname={self.botname}, model={self.model}, temperature={self.temperature}, max_tokens={self.max_tokens}, instruction={self.instruction}, tools={self.tools})'

class Thread(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    owner = db.Column(db.Integer, sa.ForeignKey('user.id'), nullable=False)
    robot = db.Column(db.Integer, sa.ForeignKey('robot.id'))
    messages = db.Column(db.Text(), nullable=False)
    embeddings = db.Column(db.Text())
    phase = db.Column(db.Enum('idle', 'responsing', 'computing', 'stopping'), default='idle')
    choices = db.Column(db.String(160), default='')
    choices_ready = db.Column(db.Boolean, default=False)
    title = db.Column(db.String(80), default='Untitled')
    title_ready = db.Column(db.Boolean, default=False)

    def __repr__(self):
        return f'Thread(id={self.id}, owner={self.owner}, robot={self.robot}, messages={self.messages}, phase={self.phase}, choices={self.choices}, title={self.title})'

def markdown(text):
    if not text:
        return '...'
    return parse_md_page(text)

@app.before_request
def before_request():
    if session.get('login'):
        return None
    if request.remote_addr and (request.remote_addr == '127.0.0.1' or request.remote_addr.startswith('192.168.') or request.remote_addr.startswith('10.0.')):
        if not session.get('login'):
            session['login'] = 1
        return None
    if request.endpoint in ('login', 'register', 'logout'):
        return None
    if request.path.startswith('/static/') or request.path == '/favicon.ico':
        return None
    return redirect(url_for('login'))

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form['username']
        if username == 'admin':
            return render_template('login.html', error='用户名 "admin" 不允许登录')
        password = request.form['password']
        if not username or not password:
            return render_template('login.html', error='用户名或密码为空')
        user = db.session.execute(sa.select(User).filter(sa.and_(sa.or_(User.username == username, User.email == username), User.password == password))).scalar()
        if user:
            session['login'] = user.id
            return redirect(url_for('index'))
        else:
            return render_template('login.html', error='用户名或密码错误')
    else:
        return render_template('login.html', error=None)

@app.route('/send_code', methods=['POST'])
def send_code():
    data = request.get_json()
    email = str(data["email"]).replace('\n', ' ')
    if not email or not re.match(r'^[^@]+@[^@]+$', email):
        return jsonify({'status': 'ERROR', 'error': '邮箱格式不正确'})
    secret_key = app.secret_key or b'='
    if isinstance(secret_key, str):
        secret_key = secret_key.encode()
    rng = random.Random(secret_key + time.time_ns().to_bytes(8) + random.randint(1, 65535).to_bytes(2))
    code = str(rng.randint(100000, 999999))
    # send the code via email
    # debug only below
    print('Verification code:', code)
    session['code'] = email + '\n' + code
    return jsonify({'status': 'OK'})

@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form['username']
        if username == 'admin':
            return render_template('register.html', error='用户名 "admin" 不允许注册')
        password = request.form['password']
        if not username or not password:
            return render_template('register.html', error='用户名或密码为空')
        email = request.form['email'].replace('\n', ' ')
        if not email or not re.match(r'^[^@]+@[^@]+$', email):
            return render_template('register.html', error='邮箱格式不正确')
        code = request.form.get('code', '')
        real_code = session.get('code')
        if not real_code:
            return render_template('register.html', error='注册验证码过期，请重新获取')
        real_email, real_code = real_code.split('\n')
        if code != real_code or email != real_email:
            return render_template('register.html', error='验证码错误')
        user = User(username=username, password=password, email=email)
        db.session.add(user)
        db.session.commit()
        user = db.session.execute(sa.select(User).filter(sa.and_(User.username == username, User.password == password))).scalar()
        if user:
            session['login'] = user.id
            return redirect(url_for('index'))
        else:
            return render_template('register.html', error='用户注册失败')
    else:
        return render_template('register.html', error=None)

@app.route('/logout')
def logout():
    session.clear()
    return redirect(url_for('login'))

@app.route('/')
def index():
    user = session['login']
    thread = request.args.get("thread", type=int, default=None)
    if thread is None:
        messages = []
        phase = 'idle'
        choices = get_option("initial_choices", [])
        choices_ready = True
        robot = request.args.get("robot", type=int, default=None)
        robot_data = db.session.execute(sa.select(Robot).filter(sa.and_(Robot.id == robot, Robot.owner == user))).scalar()
        if not robot_data:
            robot = None
    else:
        thread_data = db.session.execute(sa.select(Thread).filter(sa.and_(Thread.id == thread, Thread.owner == user))).scalar()
        if not thread_data:
            return redirect(url_for('index'))
        messages = json.loads(thread_data.messages)
        phase = thread_data.phase
        choices = thread_data.choices.split('\n')
        choices_ready = thread_data.choices_ready
        robot = thread_data.robot
        robot_data = db.session.execute(sa.select(Robot).filter(sa.and_(Robot.id == robot, Robot.owner == user))).scalar()
    thread_list = db.session.execute(sa.select(Thread).filter(Thread.owner == user)).scalars()
    robot_list = db.session.execute(sa.select(Robot).filter(Robot.owner == user)).scalars()
    return render_template('index.html',
                           messages=messages,
                           phase=phase,
                           threads=thread_list,
                           robots=robot_list,
                           current_thread=thread,
                           current_robot=robot,
                           current_robot_data=robot_data,
                           options=get_options(),
                           choices=choices,
                           choices_ready=choices_ready,
                           markdown=markdown,
                           )

@app.route('/set_robot')
def set_robot():
    user = session['login']
    thread = request.args.get("thread", type=int, default=None)
    robot = request.args.get("robot", type=int, default=None)
    if thread is None:
        return redirect(url_for('index', robot=robot))
    else:
        thread_data = db.session.execute(sa.select(Thread).filter(sa.and_(Thread.id == thread, Thread.owner == user))).scalar()
        if not thread_data:
            return redirect(url_for('index'))
        thread_data.robot = robot
        db.session.add(thread_data)
        db.session.commit()
        return redirect(url_for('index', thread=thread))

@app.route('/robot', methods=['GET', 'POST'])
def robot():
    user = session['login']
    robot = request.args.get("robot", type=int, default=None)
    if robot is None:
        robot_data = Robot(owner=user, botname=random_gb2312_name())
        db.session.add(robot_data)
        db.session.commit()
        robot = robot_data.id
        return redirect(url_for('robot', robot=robot))
    robot_data = db.session.execute(sa.select(Robot).filter(sa.and_(Robot.id == robot, Robot.owner == user))).scalar()
    if not robot_data:
        return redirect(url_for('index'))
    if request.method == 'POST':
        botname = request.form.get("botname", type=str, default=None)
        if not botname:
            return render_template('robot.html', robot=robot_data, error='未填写机器人名称')
        tools = request.form.get("tools", type=str, default=None)
        model = request.form.get("model", type=str, default=None)
        instruction = request.form.get("instruction", type=str, default=None)
        introduction = request.form.get("introduction", type=str, default=None)
        temperature = request.form.get("temperature", type=float, default=0.5)
        if not 0 <= temperature <= 1:
            return render_template('robot.html', robot=robot_data, error='温度值需要在 0~1 范围内')
        max_tokens = request.form.get("max_tokens", type=int, default=4096)
        if not 1 <= max_tokens <= 4096:
            return render_template('robot.html', robot=robot_data, error='最大生成长度需要在 1~4096 范围内')
        robot_data.tools = tools
        robot_data.instruction = instruction
        robot_data.introduction = introduction
        robot_data.model = model
        robot_data.botname = botname
        robot_data.temperature = temperature
        robot_data.max_tokens = max_tokens
        db.session.add(robot_data)
        db.session.commit()
        return redirect(url_for('index', robot=robot))
    return render_template('robot.html', robot=robot_data, error=None)

def generate_choices(thread_data: Thread):
    try:
        messages = json.loads(thread_data.messages)
        choices = conversation_choices(messages)
        thread_data.choices = '\n'.join(choices)
    except:
        traceback.print_exc()
    thread_data.choices_ready = True

def generate_title(thread_data: Thread):
    try:
        messages = json.loads(thread_data.messages)
        title = conversation_title(messages)
        thread_data.title = title
    except:
        traceback.print_exc()
    thread_data.title_ready = True

def check_safety(name, args):
    timeout = get_safety_timeout(name, args)
    if timeout <= 0:
        return True
    else:
        print('safety check:', name, json.dumps(args, ensure_ascii=False, indent=4))
        time.sleep(timeout)
        return True

def resolve_query(thread_data: Thread, robot: Robot | None):
    need_title = not thread_data.title_ready
    try:
        t0 = time.time()
        messages, nomits = auto_embed_messages(thread_data)
        dt = time.time() - t0
        pretty_print_embedding(len(messages), len(messages) - nomits, dt)
        t1 = t0
        def on_message(message, dt, append, finish):
            nonlocal t1
            t2 = time.time()
            if t2 - t1 >= 0.2 or append or finish:
                t1 = t2
                messages = json.loads(thread_data.messages)
                if append:
                    messages.append(message)
                else:
                    messages[-1] = message
                thread_data.messages = json.dumps(messages, ensure_ascii=False, separators=(',', ':'))
                db.session.add(thread_data)
                stop = thread_data.phase == 'stopping'
                db.session.commit()
                if stop:
                    raise StopIteration
                if finish:
                    pretty_print_message(message, dt, append)
        callbacks = {
            'message': on_message,
            'answer': play_audio_text,
            'check': check_safety,
        }
        if robot:
            model = robot.model or "gpt-4o"
            options = {
                'streamed': get_option('streamed', True),
                'max_tokens': max(1, min(4096, robot.max_tokens)),
                'temperature': max(0.0, min(1.0, robot.temperature)),
                'tool_categories': robot.tools.split(' ') if robot.tools else [],
            }
            if robot.instruction:
                messages.insert(0, {'role': 'system', 'content': robot.instruction})
        else:
            model = "gpt-4o"
            options = {
                'streamed': get_option('streamed', True),
                'max_tokens': 4096,
                'temperature': 0.5,
                'tool_categories': get_option('tool_categories', []),
            }
        model = get_option('available_models').get(model, model)
        ask_question(model, messages, options, callbacks)
    except StopIteration:
        thread_data.phase = 'idle'
        thread_data.choices_ready = True
        thread_data.choices = '继续说下去'
        db.session.add(thread_data)
        db.session.commit()
        return
    except:
        traceback.print_exc()
        thread_data.phase = 'idle'
        thread_data.choices_ready = True
        thread_data.choices = '继续说下去'
        db.session.add(thread_data)
        db.session.commit()
        return
    thread_data.phase = 'idle'
    db.session.add(thread_data)
    db.session.commit()
    generate_choices(thread_data)
    db.session.add(thread_data)
    db.session.commit()
    if need_title:
        generate_title(thread_data)
        db.session.add(thread_data)
        db.session.commit()

@app.route('/send_message', methods=['POST'])
def send_message():
    user = session['login']
    query = request.form['query']
    if not query:
        return redirect(url_for('index'))
    thread = request.form.get("thread", type=int, default=None)
    robot = request.form.get('robot', type=int, default=None)
    if thread is None:
        messages = []
        messages.append({'role': 'user', 'content': query})
        thread_data = Thread(owner=user, robot=robot, title=query[:40])
        thread_data.phase = 'responsing'
        thread_data.choices_ready = False
        thread_data.choices = ''
        thread_data.messages = json.dumps(messages, ensure_ascii=False, separators=(',', ':'))
        db.session.add(thread_data)
        db.session.commit()
        thread = thread_data.id
        return redirect(url_for('index', thread=thread))
    else:
        thread_data = db.session.execute(sa.select(Thread).filter(sa.and_(Thread.id == thread, Thread.owner == user))).scalar()
        if not thread_data:
            return redirect(url_for('index'))
        if thread_data.phase == 'idle':
            messages = json.loads(thread_data.messages)
            messages.append({'role': 'user', 'content': query})
            thread_data.phase = 'responsing'
            thread_data.choices_ready = False
            thread_data.choices = ''
            thread_data.messages = json.dumps(messages, ensure_ascii=False, separators=(',', ':'))
            robot_data = db.session.execute(sa.select(Robot).filter(sa.and_(Robot.id == robot, Robot.owner == user))).scalar()
            if robot_data:
                thread_data.robot = robot
            db.session.add(thread_data)
            db.session.commit()
        return redirect(url_for('index', thread=thread))

@app.route('/stop_generating', methods=['POST'])
def stop_generating():
    user = session['login']
    thread = request.form.get("thread", type=int, default=None)
    if thread is None:
        return redirect(url_for('index'))
    thread_data = db.session.execute(sa.select(Thread).filter(sa.and_(Thread.id == thread, Thread.owner == user))).scalar()
    if not thread_data:
        return redirect(url_for('index'))
    if thread_data.phase not in ('idle', 'stopping'):
        thread_data.phase = 'stopping'
        db.session.add(thread_data)
        db.session.commit()
    elif thread_data.phase == 'stopping':
        thread_data.phase = 'idle'
        db.session.add(thread_data)
        db.session.commit()
    return redirect(url_for('index', thread=thread))

@app.route('/start_response', methods=['POST'])
def start_response():
    user = session['login']
    data = request.get_json()
    thread = data["thread"]
    thread_data = db.session.execute(sa.select(Thread).filter(sa.and_(Thread.id == thread, Thread.owner == user))).scalar()
    if not thread_data:
        return jsonify({'status': 'ERROR', 'error': 'invalid thread id'}), 400
    robot = None
    if thread_data.robot:
        robot = db.session.execute(sa.select(Robot).filter(sa.and_(Robot.id == thread_data.robot))).scalar()
    if thread_data.phase == 'responsing':
        thread_data.phase = 'computing'
        db.session.add(thread_data)
        db.session.commit()
        resolve_query(thread_data, robot)
        return jsonify({'status': 'OK'})
    if thread_data.phase in ('computing', 'stopping'):
        return jsonify({'status': 'ERROR', 'error': 'thread busy'})
    else:
        return jsonify({'status': 'ERROR', 'error': 'thread idle'})

@app.route('/get_response', methods=['POST'])
def get_response():
    user = session['login']
    data = request.get_json()
    thread = data["thread"]
    thread_data = db.session.execute(sa.select(Thread).filter(sa.and_(Thread.id == thread, Thread.owner == user))).scalar()
    data = {'responsing': thread_data.phase != 'idle', 'messages': json.loads(thread_data.messages)}
    return jsonify(data)

@app.route('/get_title', methods=['POST'])
def get_title():
    user = session['login']
    data = request.get_json()
    thread = data["thread"]
    thread_data = db.session.execute(sa.select(Thread).filter(sa.and_(Thread.id == thread, Thread.owner == user))).scalar()
    data = {'title_ready': thread_data.title_ready, 'title': thread_data.title}
    return jsonify(data)

@app.route('/get_choices', methods=['POST'])
def get_choices():
    user = session['login']
    data = request.get_json()
    thread = data["thread"]
    thread_data = db.session.execute(sa.select(Thread).filter(sa.and_(Thread.id == thread, Thread.owner == user))).scalar()
    data = {'choices_ready': thread_data.choices_ready, 'choices': thread_data.choices.split('\n')}
    return jsonify(data)

@app.route('/favicon.ico')
def favicon():
    return redirect(url_for('static', filename='favicon.ico'))

def init_db():
    db.create_all()
    admin = db.session.execute(sa.select(User).filter(User.username == 'admin')).scalar()
    if admin is None:
        password = b64encode(os.urandom(16)).strip(b'=').decode()
        print('Admin password: ' + password)
        admin = User(id=1, username='admin', password=password, email='admin', admin=True)
        db.session.add(admin)
        db.session.commit()
    elif admin.id != 1:
        raise ValueError('Admin user id must be 1')

if __name__ == '__main__':
    with app.app_context():
        db.drop_all()
        init_db()
    app.run(debug=True)
